Skip to content

Hold MLOps artifacts on Unity Catalog volumes through the Files API - #114

Merged
Vedaant-J merged 25 commits into
mainfrom
vedaant/mlops-dbx-volume
Nov 3, 2026
Merged

Vedaant-J merged 25 commits into
mainfrom
vedaant/mlops-dbx-volume

Conversation

@Vedaant-J

@Vedaant-J Vedaant-J commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor

Stack: main (#53 merged) → #112 → #114 → #50 → #60 → #64 → #115. Sits between #112 and #50. Merge bottom to top.

Summary

A /Volumes/<catalog>/<schema>/<volume>/... path is reachable only through the FUSE mount on Databricks compute, so a training job on a laptop, EC2 box or CI cannot publish a model or export scores to the volume a Databricks job reads. This PR makes such a path a file destination of DatabricksConnector from any machine. The sink implements #112's FileStoreSink and reads and writes the volume through the Databricks Files API with the connector's identity, reusing #112's writers. Files up to 4 GiB are plain objects that Spark reads directly; larger ones are stored as 4 GiB parts (the Files API caps one upload at 5 GiB). No SQL warehouse is used.

Usage

root = "/Volumes/main/kumo/artifacts"
trainer.fit(training_table, snapshot=snapshot, export_config=ModelOutputConfig(
    databricks, path=f"{root}/models/churn"))
trainer = Trainer.load_from_connector(databricks, path=f"{root}/models/churn")
trainer.predict(prediction_table, snapshot=snapshot, output_config=PredictionOutputConfig(
    databricks, output_paths={"predictions": f"{root}/scores/"}))

The identity needs USE CATALOG, USE SCHEMA, READ VOLUME (load) and WRITE VOLUME (export).

Back-compat

Situation Behaviour
Table name on a DatabricksConnector Unchanged: COPY INTO through the staging volume.
Unity Catalog registry publish and fetch (location=) Unchanged.
/Volumes/... path on a DatabricksConnector Was refused; now written through the Files API.
/Volumes/... string root on Databricks compute Unchanged: a FileConnector over FUSE.
Export into a volume directory that already holds files Fails before scoring starts: Output directory '/Volumes/.../scores/' already contains data from another export. Choose an empty destination path.
Upload or download failure FileNotFoundError, PermissionError or OSError with the volume path; no signed URL in the message. A failed write cancels its upload, so nothing is committed.
A file larger than 4 GiB read by Spark A directory of <name>.chunks/<generation>/ parts; only kumo reads it as one file.

Testing

  • kumo-databricks/tests/test_volume.py (10 tests, in-memory Files API in tests/fakes.py): reads across part boundaries; sized writes plain up to a part; failed and refused uploads; a download that ignores Range is refused; directories and missing paths; bundle publish and fetch, Parquet and frame exports through a volume path; an occupied directory refused before scoring.
  • Live with the rest of the series applied, on a test workspace with synthetic data, off-cluster and in a cluster job: 3.0 GiB and 4.29 GiB files round-trip with matching CRCs, a read across the 4 GiB part boundary is exact, and a kumo.mlops fit and predict with the artifact root on a volume writes snapshots, models and predictions there.

Change map: see the kumo-pr-change-map comment.

🤖 Generated with Claude Code

@Vedaant-J Vedaant-J self-assigned this Oct 24, 2026
@copy-pr-bot

copy-pr-bot Bot commented Oct 24, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 24, 2026 19:25 — with GitHub Actions Inactive
@coderabbitai

coderabbitai Bot commented Oct 24, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • Resume reviews
  • Trigger review
Walkthrough

The Databricks connector now routes Unity Catalog volume paths through a Files API-backed filesystem. The filesystem supports metadata lookup, listing, deletion, ranged reads, and streamed uploads. The sink backend classifies volume paths and provides the filesystem for those paths while retaining table-writing behavior. Tests cover file operations, model bundle and prediction flows, and destination preflight checks. The README documents volume path usage, permissions, and file-size handling.

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant DatabricksSinkBackend
  participant VolumeHandler
  participant FilesAPI
  Caller->>DatabricksSinkBackend: Request filesystem for volume path
  DatabricksSinkBackend->>VolumeHandler: Route file operation
  VolumeHandler->>FilesAPI: Upload, download, list, or delete
  FilesAPI-->>VolumeHandler: Return file data or operation result
  VolumeHandler-->>Caller: Return filesystem result
Loading
Merge Risk: Moderate · previous review 2bc33

Large CSV exports to Unity Catalog volumes can use over 1 GiB of memory during upload and may terminate memory-constrained jobs. Bound the SDK upload buffering before merging. The README also links to an MLOps doc that does not exist until a later PR in the stack lands.

Important

Pre-merge checks failed

Please resolve all errors before merging. Addressing warnings is optional.

Failed Failed checks: 1 error · 1 inconclusive

Check name Status Explanation Resolution
Public Resource Access Error The added required repository reference packages/kumo-engine/docs/mlops.md at packages/kumo-databricks/README.md:87 is inaccessible. The path is absent from both reviewed base and head trees, and … Replace the reference with an existing public documentation path, or add packages/kumo-engine/docs/mlops.md to the repository and link to that file.
Change Map Design Review Inconclusive The required PR conversation comment beginning <!-- kumo-pr-change-map --> is not available in the supplied review context or repository checkout. The description only links to the map and the objec… Provide the exact current conversation comment containing the change map. Include its rationale, observed failure or release requirement, existing behavior, smallest sufficient change, focused verification, stopping condition, and treatment…
Passed Passed checks: 1 passed
Check name Status Explanation
Pr Size Passed The pull request changes 6 files. This does not exceed the failure threshold of 40 files.
Full details: Public Resource Access

Explanation

The added required repository reference packages/kumo-engine/docs/mlops.md at packages/kumo-databricks/README.md:87 is inaccessible. The path is absent from both reviewed base and head trees, and the anonymous public URL https://raw.githubusercontent.com/NVIDIA/kumo-engine/main/packages/kumo-engine/docs/mlops.md returns HTTP 404.

Resolution

Replace the reference with an existing public documentation path, or add packages/kumo-engine/docs/mlops.md to the repository and link to that file.

Full details: Change Map Design Review

Explanation

The required PR conversation comment beginning &lt;!-- kumo-pr-change-map --&gt; is not available in the supplied review context or repository checkout. The description only links to the map and the objective summary paraphrases it, so I cannot verify the required per-change rationale, existing behavior, smallest sufficient change, focused verification, or stopping condition in that same comment. Direct code evidence also shows an unresolved design question: volume.py:87-90 calls files.upload without upload-buffer controls, and plain_bytes is set to 4 GiB; the existing tests cover only small in-memory uploads. The missing map prevents deciding whether this risk and the required design evidence were addressed.

Resolution

Provide the exact current conversation comment containing the change map. Include its rationale, observed failure or release requirement, existing behavior, smallest sufficient change, focused verification, stopping condition, and treatment of the plain-object upload buffering risk.

  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

CodeRabbit Verify

  • Run CodeRabbit Verify to test the changes in this pull request.
Recent review info
Run configuration
Commits

Reviewing files that changed from the base of the PR and between 5fb2a8d and 2bc3351.

Files selected for processing: 3
  • packages/kumo-databricks/README.md
  • packages/kumo-databricks/src/kumo_databricks/volume.py
  • packages/kumo-databricks/tests/test_volume.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

@Vedaant-J

Vedaant-J commented Oct 24, 2026 •

Copy link
Copy Markdown
Contributor Author

Change map

File Function / scope What changed Why
packages/kumo-databricks/src/kumo_databricks/volume.py new: VolumeHandler, volume_filesystem, is_volume_path Chunked filesystem over the Files API; 4 GiB parts, ranged reads Volumes cap one upload at 5 GiB
packages/kumo-databricks/src/kumo_databricks/sink.py _DatabricksTables, DatabricksSinkBackend: holds_files, filesystem Table methods split out; FileStoreRouting sends volume paths through FilesystemSink Exports and bundles on /Volumes/... paths
packages/kumo-databricks/src/kumo_databricks/registry.py _address Error names the volume path form Points users at the supported path
packages/kumo-databricks/README.md "Artifacts on a volume" section Usage, file layout, required grants Setup the code cannot show
packages/kumo-databricks/tests/fakes.py FakeVolume, FakeMetadata, FakeStreamingResponse In-memory volume with ranged downloads Tests run without a workspace
packages/kumo-databricks/tests/test_volume.py new tests Parts, failures, permissions, bundles, exports, occupied directory Covers the volume sink

@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from cab9cb7 to 91d494f Compare October 24, 2026 20:31
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 24, 2026 20:31 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 24, 2026 20:43 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from 044ad8c to b322307 Compare October 26, 2026 18:28
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 26, 2026 18:29 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J changed the base branch from vedaant/mlops-prediction-tracking to vedaant/mlops-filestore-sink October 26, 2026 18:29
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from b322307 to d1d6d05 Compare October 27, 2026 07:30
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 27, 2026 07:30 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from d1d6d05 to 3f3d201 Compare October 27, 2026 09:51
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 27, 2026 09:51 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from 3f3d201 to 98fb53e Compare October 27, 2026 18:58
@Vedaant-J
Vedaant-J temporarily deployed to gha-kumo-workflow-tests October 27, 2026 18:59 — with GitHub Actions Inactive
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from 98fb53e to e2d71ad Compare October 28, 2026 16:26
Vedaant-J and others added 25 commits November 3, 2026 11:49
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…FilesystemSink

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…equality

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…tream log

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ep only table writes

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
CSV exports open the stream without a size, so anything over 256 MiB
rolled over to .chunks parts that Spark cannot read as one file.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The export preflight lists the destination before writing.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Pass an explicit 16 MiB part size and 4 upload threads to files.upload
when the SDK accepts them (databricks-sdk 0.70+), so a large plain object
no longer buffers up to parallelism x part_size with the SDK's defaults.
Older SDKs upload single-threaded in 10 MiB chunks and get no extra args.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@Vedaant-J
Vedaant-J force-pushed the vedaant/mlops-dbx-volume branch from fb03edd to 01ca25a Compare November 3, 2026 19:49
@Vedaant-J
Vedaant-J deployed to gha-kumo-workflow-tests November 3, 2026 19:49 — with GitHub Actions Active
@Vedaant-J
Vedaant-J merged commit 5f0f6f4 into main Nov 3, 2026
46 checks passed
@Vedaant-J
Vedaant-J deleted the vedaant/mlops-dbx-volume branch November 3, 2026 19:49
Vedaant-J added a commit that referenced this pull request Nov 3, 2026
**Stack:** main (#53 merged) → #112 → #114 → **#50** → #60 → #64 → #115.
Merge bottom to top.

## Summary

Every graph snapshot, training table and model bundle now carries a
self-contained `_provenance.json` that leads back to the source table
versions (and locations, wherever the source backend can name them),
query, plan and engine build that produced it. Today a model has no id,
an artifact records no lineage, an index works only on the machine that
built it, and scoring needs the training job's `Graph`.

This PR adds `kumo.MLOps` / `with kumo.mlops(...)` to store each
artifact once under one artifact root (`s3://`, `gs://`, local, or a
`(connector, path)` file store from #112); `GraphSnapshot.export` /
`load_from_connector` to move a checksummed index between machines;
`snapshot.source_changes(graph)` to see which tables changed since a
snapshot; and `Trainer.predict(snapshot)` to score a stored snapshot
with the trained query. `MLOps(on_error='raise')` turns a failed store
into a `RuntimeError` after the call finishes. Off unless asked for.

**Design decision:** for this release MLOps is opt-in per call or per
block (`mlops=` or `with kumo.mlops(...)`). A call outside both runs
without MLOps, and `on_error="raise"` does not detect that. A model- or
workflow-bound association with an explicit `mlops=False` was considered
and is left as a follow-up. See "Per call, not per model" in
`docs/mlops.md`.

## Usage

```python
cfg = kumo.MLOps("s3://<bucket>/kumo")             # store=True: snapshots, tables, models
with kumo.mlops(cfg):                               # or mlops=cfg on one call
    snapshot = graph.snapshot()
    result = trainer.fit(pq.generate_training_table(snapshot), snapshot=snapshot)

# Later, on another machine with no Graph code:
snapshot = GraphSnapshot.load_from_connector(cfg, id=result.snapshot_id)
trainer = Trainer.load_from_connector(cfg, path=f"models/{result.id}")
scores = trainer.predict(snapshot)
```

## Back-compat

Without a block or `mlops=`, existing calls behave as before.

| Situation | Behaviour |
|---|---|
| Index built before this PR | Reused, not rebuilt; its manifest is
backfilled with `"reconstructed": true`. |
| Bundle published before this PR | Loads and predicts identically; id
computed at load, `manifest` is `None`. |
| Artifact from this PR read by an older engine | Extra files are
ignored. |
| A single S3 or GCS object source | Versioned by VersionId or entity
tag, so its snapshot gets a new id once. |
| `verify_model_bundle` return value | 64 hex instead of a 16-hex digest
no caller used. |
| `DatabricksConnector(identity=...)` | Unchanged; `identity` is now
optional (ambient auth). |
| Manifest, sidecar or root write fails | A warning; the call completes
as before. A failed store raises a `RuntimeError` instead, once the call
has finished, with `MLOps(on_error='raise')`. |

## Testing

- kumo-engine: `test_snapshot_export.py` (round trip, reuse, interrupted
and concurrent exports, corrupt markers, `source_changes`, `on_error`),
`test_snapshot_manifest.py`, `test_training_table_export.py`,
`test_holdout_df.py` (published model's holdout, loaded and fetched
holdout, empty test split), `test_bundle_provenance.py`,
`test_query_from_snapshot.py`, `test_filestore_sink.py` (identical
bundle reused), `test_mlops.py`, `test_graph_mismatch.py`,
`test_artifacts.py`.
- kumo-core typed source versions, bundle ids and export sidecars;
kumo-s3 `test_sink_delete.py`; Snowflake, Databricks and BigQuery
identity tests assert `version_kind`.
- DiskGraph: Rust
`range_checksums_cover_every_listed_file_and_the_manifest`; pybind
`test_exception_exports.py`.
- Env-gated live S3 and GCS round trips in
`test_snapshot_export_live.py`.

Change map: see the `kumo-pr-change-map` comment.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

This branch was successfully deployed

1 active deployment
gha-kumo-workflow-tests — 01ca25a4 Deployed Nov 3, 2026 by Vedaant-J via import #1375
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants